
#include <tftpserver.hpp>
#include <sys/socket.h>
#include <sys/time.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <string.h>
#include <unistd.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <dirent.h>
#include <stdio.h>
#include <string>
#include <functional>
#include <queue>
#include <map>

extern void print_log_stamp(const char * pFuntion, const char* pModule, const int line,const char *format, ...);
#define MODULE_NAME  "tftp"
#define LOGTIME(...) print_log_stamp(__FUNCTION__,MODULE_NAME,__LINE__,__VA_ARGS__)

/**<   通过执行命令`export tftpserver_log_enable=yes`可以使能LOG */
#define SHOW_LOG_ENV_STR "tftpserver_log_enable"
static bool __logflg = 
    (nullptr != getenv(SHOW_LOG_ENV_STR))?
    (!strncmp(getenv(SHOW_LOG_ENV_STR),"yes",3)?
        true:false):
    false;

#define LOG(msg...) do{\
    if(__logflg){\
        LOGTIME(msg);}\
    }while(0)

#define ERRLOG(msg...) do{\
        LOGTIME(msg);\
    }while(0)

#define FILE_ASCII "netasci"
#define FILE_BINARY "octet"

using namespace std;

static char *tftp_err_str(uint32_t id)
{
    const char *str[]=
    {
        "Not defined",/**<  tftp_err_undefined */
        "File not found", /**<  tftp_err_not_found */
        "Access violation", /**<  tftp_err_access_violation */
        "Disk full or allocation exceeded", /**<  tftp_err_disk_full */
        "Illegal TFTP operation", /**<  tftp_err_illegal_operation */
        "Unknown transfer ID", /**<  tftp_err_unknown_transfer_id */
        "File already exists", /**<  tftp_err_file_exist */
        "No such user", /**<  tftp_err_no_such_user */
    };

    if(id>=8)
    {
        return (char*)str[0];
    }
    return (char*)str[id];
}

static std::map<int16_t,tftp_agent*> workagent; /**<  正在运行的agent链表 */
void tftp_agent::tftp_demon_loop(void)
{
    #define DEACTIVE_ALLOW  (60)
    #define DEMON_CHECK_PERIOD (2)
    LOG("");
    while(1)
    {
        std::this_thread::sleep_for(std::chrono::seconds(DEMON_CHECK_PERIOD));
        uint32_t now = time((time_t*)nullptr);

        for(auto it  = workagent.begin();
                 it != workagent.end();
                 it++)
        {
            /**< 删除长时间不活动的客户端 */
            if(now - it->second->stamp >= DEACTIVE_ALLOW)
            {
                LOG("client(port=%d) deactive too long(more than %s seconds)",
                    it->second->port,DEACTIVE_ALLOW);
                it->second->running = false;
                workagent.erase(it);
                it->second->~tftp_agent();
            }
        }
    }
}

static std::queue<uint16_t> freeport; /**<  所有可用的端口号都放在这个队列中 */
static std::once_flag oc;

tftp_agent:: tftp_agent(const tftp_request_packet_s req,const int sock,struct sockaddr_in clt)
{
    LOG("");

    std::call_once(oc,[&]()
    {
        LOG("initial free port queue ...");
        for(int i=0;i<TFTP_AGENT_NUM;i++)
        {
            freeport.push(TFTP_AGENT_START+i);
        }
        workagent.clear();
        std::thread demon(&tftp_agent::tftp_demon_loop);
        demon.detach();
    });
    if(0>sock)
    {
        /**<  非法的套接字一律拒绝 */
        ERRLOG("error sock");
        return;
    }

    uint16_t cmd = ntohs(req.opcode);
    if((tftp_cmd_rrq != cmd) && (tftp_cmd_wrq != cmd))
    {
        /**<  非读写请求一律拒绝 */
        ERRLOG("unsupported tftp request(%d)",req.opcode);
        return;
    }

    //this->sock = sock;
    int flag = O_RDONLY;
    if(tftp_cmd_wrq == cmd)
    {
        flag = O_RDWR | O_CREAT | O_TRUNC; 
    }
    
    /**<  解码出文件名,模式,文件大小,, */
    uint32_t index = 0;
    string fname = req.filemsg;
    index += fname.length()+1;
    string fmode = (char*)(&req.filemsg[index]);
    index += (fmode.length()+1);
    string fsizemark = (char*)(&req.filemsg[index]);
    index += (fsizemark.length()+1);

    /**<  登记好客户端 */
    client  = clt;
    cltlen  = sizeof(client);
    /**<  获得本地文件的句柄 */
    LOG("file message {name:%s,mode:%s}",fname.c_str(),fmode.c_str());
    fd = open(fname.c_str(),flag);
    if(0>fd)
    {
        LOG("open file %s fail (ret=%d)",fname.c_str(),fd);
        error_ack((char*)"file not exist?");
        return;
    }


    /**<  绑定到一个新的端口(原端口空闲出来给新的客户端使用) */
    if(freeport.empty())
    {
        ERRLOG("free port have used up!");
        error_ack(sock,(char*)"free port have used up!");
        close(fd);
        return;
    }

    port = freeport.front();
    int newsock =socket(AF_INET, SOCK_DGRAM, 0);
    if(0 > newsock)
    {
        printf("Server socket could not be created(ret=%d)\n",newsock);
        perror("socket ");
        close(fd);
		return;
	}
    struct sockaddr_in svr;
    uint32_t svrlen = sizeof(svr);
    bzero(&svr,sizeof(svr));
    svr.sin_family = AF_INET;
	svr.sin_addr.s_addr = INADDR_ANY;
	svr.sin_port = htons(port);
    int ret;
    ret = bind(newsock,(struct sockaddr*)&svr,svrlen);
	if(0 > ret)
    {
        printf("bind to port %d fail(ret=%d)\n",port,ret);
        error_ack(sock,(char*)"bind to new port fail!");
		return ;
	}
    this->sock = newsock;
    freeport.pop(); /**<  已经使用的端口从队列中删除 */
    LOG("there are %d free port",freeport.size());

    running = true;
    offset  = 0;
    /**<  成功获取端口的agent都置于监控中 */
    workagent.insert(std::pair<uint16_t,tftp_agent*>(port,this));

    if(tftp_cmd_rrq == cmd)
    {
        /**<  获取本地文件的大小 */
        struct stat statbuf={0}; 
        stat(fname.c_str(),&statbuf);
        tsize  = statbuf.st_size;
        LOG("create read loop sock=%d,newsock=%d,filefd=%d,name=%s,tsize=%d",
            sock,this->sock,fd,fname.c_str(),tsize);
        this->thd = new std::thread(
            std::bind(&tftp_agent::tftp_svr_read_loop,this));
    }
    else if(tftp_cmd_wrq == cmd)
    {
        string fsize = (char*)(&req.filemsg[index]);
        tsize = atoi(fsize.c_str()); /**<  记录文件的大小 */
        LOG("create write loop for sock=%d,newsock=%d,filefd=%d,name=%s,tsize=%d",
            sock,this->sock,fd,fname.c_str(),tsize);
        this->thd = new std::thread(
            std::bind(&tftp_agent::tftp_svr_write_loop,this));
    }
    this->thd->detach();
}

void tftp_agent:: tftp_svr_write_loop(void)
{
    LOG("thread_id=%u",pthread_self());

    /**<  接收到对方的WRQ之后,回复一个rrq/wrq的回应包 */
    if(!file_ack())
    {
        ERRLOG("file_ack fail");
        running = false;
    }

    /**<  循环执行"接收数据-写文件-发送回答"的操作步骤 */
    uint32_t cnt = 0;
    while(running)
    {
        stamp = time((time_t*)nullptr); /**<  记录当前时间 */
        /**<  接收DATA数据包,使用临时sockaddr */
        if(!data_recv(cnt))
        {
            running = false;
            break;
        }
        
        /**<  写文件 */
        int ret = write(fd,datapkt.payload,cnt);
        if(0 > ret)
        {
            ERRLOG("write %d bytes to file fail",cnt);
            error_ack((char*)"write file fail");
            break;
        }
        LOG("write %d bytes",ret);

        offset += ret;
        data_ack(tftp_cmd_ack,datapkt.block);
        if(offset == tsize)
        {
            /**<  全部数据传输完成 */
            LOG("data transmit finished");
            running = false;
            break;
        }
    }

    ERRLOG("close write server ....");
    /**<  由析构函数完成文件句柄及网络句柄的释放 */
    this->~tftp_agent();
}

void tftp_agent:: tftp_svr_read_loop(void)
{
    LOG("thread_id=%u",std::this_thread::get_id());

    /**<  接收到对方的rrq之后,回复一个rrq/wrq的回应包 */
    if(!file_ack())
    {
        ERRLOG("file_ack fail");
        running = false;
    }

    /**<  循环执行"等待回答-读文件-发送数据"的操作步骤 */
    int block = 0;
    while(running)
    {
        stamp = time((time_t*)nullptr); /**<  记录当前时间 */
        if(!ack_recv())
        {
            running = false;
            break;
        }
        if(tftp_cmd_ack != ackpkt.opcode
              || block  != ackpkt.block)
        {
            /**<  接收到的ACK与发送的数据不一致! */
            ERRLOG("ack packet is error!(cmd=%d,block=%d)",ackpkt.opcode,ackpkt.block);
            error_ack((char*)"ack packet is err");
            running = false;
            break;
        }
        block++;
        //
        memset(datapkt.payload,0,TFTP_PAYLOAD_MAX);
        int bytes = read(fd,datapkt.payload,TFTP_PAYLOAD_MAX);
        if(0 > bytes)
        {
            ERRLOG("read data from file fail(ret=%d)",bytes);
            error_ack((char*)"read file fail");
            running = false;
            break;
        }
        else if(0 == bytes)
        {
            /**<  文件全部数据读完,结束操作 */
            running = false;
            break;
        }
        
        /**<  发送DATA数据包,使用构造时的sockaddr  */
        if(!data_send(block,datapkt,bytes))
        {
            /**<  多次重发失败,立即退出 */
            ERRLOG("send timeout,quit...");
            running = false;
            break;
        }

        offset += bytes;
        if(offset == tsize)
        {
            /**<  全部数据传输完毕,退出循环发送 */
            LOG("data transmit finished");
            running = false;
            break;
        }
    }
    ERRLOG("close read server ....");
    /**<  由析构函数完成文件句柄及网络句柄的释放 */
    this->~tftp_agent();
}

tftp_agent:: ~tftp_agent()
{
    LOG("");
    if(0 < sock)
    {
        close(sock);
    }
    if(0 <= fd)
    {
       close(fd); 
    }
    LOG("free port is %d",freeport.size());
    if(0 != port)
    {
        freeport.push(port);
    }
    workagent.erase(port);
}


bool tftp_agent:: data_ack(uint16_t cmd,uint16_t block)
{
    ackpkt.opcode = htons(cmd);
    ackpkt.block  = htons(block);
    int ret = sendto(sock,&ackpkt,sizeof(tftp_ack_packet_s),0,(struct sockaddr*)&client,cltlen);
    if(0 > ret)
    {
        /**<  发送ACK,使用构造时的sockaddr */
        ERRLOG("send ack fail");
        return false;
    }
    return true;
}

bool tftp_agent:: file_ack(void)
{
    /**<  填充数据内容 */
    bzero(reqack.message,TFTP_PAYLOAD_MAX);
    uint32_t payloadlen = strlen("tsize");
    reqack.opcode   = htons(tftp_cmd_req_ack);
    strcpy(&reqack.message[0],"tsize");
    char strtsize[32];
    sprintf(strtsize,"%d",tsize);
    strcpy(&reqack.message[strlen("tsize")+1],strtsize);
    payloadlen += (strlen(strtsize)+2); /**<  opcode以及数据包发出去即可,无需发送最大的数据包 */
    int ret = sendto(sock,&reqack,2+payloadlen,0,(struct sockaddr*)&client,cltlen);
    if(0 > ret)
    {
        ERRLOG("send ack fail");
        return false;
    }
    return true;
}

bool tftp_agent:: error_ack(char *msg)
{
    bzero(&errpkt,sizeof(tftp_error_packet_s));
    errpkt.opcode  = htons(tftp_cmd_err);
    errpkt.errcode = htons(0);
    if(msg)
    {
        strncpy(errpkt.message,msg,TFTP_PAYLOAD_MAX-1);
    }
    int ret = sendto(sock,&errpkt,4+strlen(errpkt.message)+1,0,(struct sockaddr*)&client,cltlen);
    if(0 > ret)
    {
        /**<  发送ACK,使用构造时的sockaddr */
        ERRLOG("send ack fail");
        return false;
    }
    return true;
}

bool tftp_agent:: error_ack(int sock,char *msg)
{
    bzero(&errpkt,sizeof(tftp_error_packet_s));
    errpkt.opcode  = htons(tftp_cmd_err);
    errpkt.errcode = htons(0);
    if(msg)
    {
        strncpy(errpkt.message,msg,TFTP_PAYLOAD_MAX-1);
    }
    int ret = sendto(sock,&errpkt,4+strlen(errpkt.message)+1,0,(struct sockaddr*)&client,cltlen);
    if(0 > ret)
    {
        /**<  发送ACK,使用构造时的sockaddr */
        ERRLOG("send ack fail");
        return false;
    }
    return true;
}

bool tftp_agent:: error_ack(uint16_t err,char *msg)
{
    bzero(&errpkt,sizeof(tftp_error_packet_s));
    errpkt.opcode  = htons(tftp_cmd_err);
    errpkt.errcode = htons(err);
    if(msg)
    {
        strncpy(errpkt.message,msg,TFTP_PAYLOAD_MAX-1);
    }
    int ret = sendto(sock,&errpkt,4+strlen(errpkt.message)+1,0,(struct sockaddr*)&client,cltlen);
    if(0 > ret)
    {
        /**<  发送ACK,使用构造时的sockaddr */
        ERRLOG("send ack fail");
        return false;
    }
    return true;
}

bool tftp_agent:: data_send(uint16_t block,tftp_data_packet_s &pkt,uint32_t bytes)
{
    #define SEND_WAIT_MAX   (200)
    #define RECV_DELAY      (100)

    pkt.opcode = htons(tftp_cmd_data);
    pkt.block  = htons(block);
    int wait_times     = SEND_WAIT_MAX;
    while(wait_times)
    {
        int ret = sendto(sock,&pkt,4+bytes,0,(struct sockaddr*)&client,cltlen);
        if(0 > ret)
        {
            ERRLOG("send fail (ret=%d)",ret);
            //可能是网络通信导致的问题,适当延时重发送
            std::this_thread::sleep_for(std::chrono::milliseconds(RECV_DELAY));
            wait_times--;
        }
        else
        {
            LOG("send %d bytes data",bytes);
            return true;
        }
    }
    LOG("send timeout");
    return false;
}

bool tftp_agent:: data_recv(uint32_t &cnt)
{
    #define RECV_DELAY      (100)
    #define RECV_WAIT_MAX   (200)

    struct sockaddr_in  addr; /**<  实际的客户端 */
    uint32_t addrlen = sizeof(addr);
    int wait_times   = RECV_WAIT_MAX;

    while(wait_times)
    {
        bzero(&datapkt,sizeof(tftp_data_packet_s));
        int ret = recvfrom(sock,&datapkt,sizeof(tftp_data_packet_s),0,(struct sockaddr*)&addr,&addrlen);
        if(0 > ret)
        {
            ERRLOG("recv fail (ret=%d),%s",ret,strerror(errno));
            //可能是网络通信导致的问题,适当延时重接收
            std::this_thread::sleep_for(std::chrono::milliseconds(RECV_DELAY));
            wait_times--;
            continue;
        }
        datapkt.block  = ntohs(datapkt.block);
        datapkt.opcode = ntohs(datapkt.opcode);
        cnt = (ret-4);
        LOG("recv %d bytes cmd=%d,block=%d,cnt=%d",ret,datapkt.opcode,datapkt.block,cnt);
        return true;
    }
    return false;
}

bool tftp_agent:: ack_recv(void)
{
    struct sockaddr_in  addr; /**<  实际的客户端 */
    uint32_t addrlen = sizeof(addr);

    int ret = recvfrom(sock,&ackpkt,sizeof(tftp_ack_packet_s),0,(struct sockaddr*)&addr,&addrlen);
    if(0 > ret)
    {
        ERRLOG("recv ack fail(ret=%d)",ret);
        return false;
    }
    ackpkt.opcode = ntohs(ackpkt.opcode);
    ackpkt.block  = ntohs(ackpkt.block);
    return true;
}